package com.ruoyi.ai.factory.impl;

import com.alibaba.dashscope.common.Role;
import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.builder.ChatBuilder;
import com.google.common.collect.Lists;
import com.ruoyi.ai.domain.vo.WxChatVo;
import com.ruoyi.ai.domain.vo.WxModelConfigVo;
import com.ruoyi.ai.enums.MfgEnum;
import com.ruoyi.ai.factory.ModelChatTemplate;
import com.ruoyi.ai.ws.ChatWebSocketManager;
import com.ruoyi.common.exception.ServiceException;
import com.ruoyi.common.utils.StringUtils;
import io.github.lnyocly.ai4j.listener.SseListener;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletion;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatMessage;
import io.github.lnyocly.ai4j.service.IChatService;
import io.github.lnyocly.ai4j.service.PlatformType;
import io.github.lnyocly.ai4j.service.factor.AiService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * Ollama模型服务
 */
@Slf4j
@Service
public class OllamaModelChatService extends ModelChatTemplate {
    @Autowired
    private AiService aiService;

    public OllamaModelChatService() {
        super(MfgEnum.OLLAMA.getType());
    }

    @Resource
    private ChatWebSocketManager chatWebSocketManager;

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId) {
        try {
            log.info("发送的请求为{}", chatVoList);
            
            // 定义消息处理函数，用于发送websocket消息
            Consumer<String> messageHandler = message -> {
                chatWebSocketManager.sendOneMessage(userId, message);
            };
            
            // 使用公共方法执行流式聊天
            String result = doChatStream(chatVoList, messageHandler);
            
            //发送对话结束标记，前端好判断本次对话已结束，允许下一轮对话
            chatWebSocketManager.sendOneMessage(userId, OVER_FLAG);
            return result;
        } catch (Exception e) {
            log.error("调用Ollama接口出现问题：{}", e.getMessage());
            throw new ServiceException("调用Ollama接口出现问题");
        }
    }

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId, String prompt) {
        if (chatVoList.size() == 1 && StringUtils.isNotBlank(prompt)) {
            WxChatVo promptVo = new WxChatVo();
            promptVo.setRole(Role.SYSTEM.getValue());
            promptVo.setContent(prompt);
            chatVoList.add(0, promptVo);
        }
        return chatStream(chatVoList, userId);
    }

    @Override
    public String chatStreamOnce(String content) {
        try {
            List<WxChatVo> chatVoList = Lists.newArrayList();
            WxChatVo wxChatVo = new WxChatVo();
            wxChatVo.setRole(Role.USER.getValue());
            wxChatVo.setContent(content);
            chatVoList.add(wxChatVo);
            
            // 使用公共方法执行流式聊天，不发送websocket消息
            return doChatStream(chatVoList, null);
        } catch (Exception e) {
            log.error("调用Ollama接口出现问题：{}", e.getMessage());
            return "调用大模型接口失败，请稍后尝试";
        }
    }

    @Override
    public String chatStreamOnce(String content, String prompt) {
        try {
            // 构建聊天内容列表
            List<WxChatVo> chatVoList = Lists.newArrayList();
            
            // 添加系统提示
            if (StringUtils.isNotBlank(prompt)) {
                WxChatVo promptVo = new WxChatVo();
                promptVo.setRole(Role.SYSTEM.getValue());
                promptVo.setContent(prompt);
                chatVoList.add(promptVo);
            }
            
            // 添加用户消息
            WxChatVo userVo = new WxChatVo();
            userVo.setRole(Role.USER.getValue());
            userVo.setContent(content);
            chatVoList.add(userVo);
            
            // 使用公共方法执行流式聊天，不发送websocket消息
            return doChatStream(chatVoList, null);
        } catch (Exception e) {
            log.error("调用Ollama接口出现问题：{}", e.getMessage());
            return "调用大模型接口失败，请稍后尝试";
        }
    }
    
    /**
     * 执行流式聊天的公共方法
     * @param chatVoList 聊天消息列表
     * @param messageHandler 消息处理函数，如果为null则不发送websocket消息
     * @return 聊天结果
     * @throws Exception 异常
     */
    private String doChatStream(List<WxChatVo> chatVoList, Consumer<String> messageHandler) throws Exception {
        // 获取chat服务实例
        IChatService chatService = aiService.getChatService(PlatformType.OLLAMA);

        //
        // 构造基本请求参数
        ChatCompletion chatCompletion = getChatCompletionBase();
        // 构造消息参数
        chatCompletion.setMessages(buildChatMessageList(chatVoList));
        
        // 构造自定义监听器
        final StringBuilder resultBuilder = new StringBuilder();
        SseListener sseListener = new SseListener() {
            @Override
            protected void send() {
                String currStr = this.getCurrStr();
                resultBuilder.append(currStr);
                
                // 如果有消息处理函数，则发送消息
                if (messageHandler != null) {
                    messageHandler.accept(currStr);
                }
            }
        };
        // 显示函数参数，默认不显示
        sseListener.setShowToolArgs(true);

        // 发送SSE请求
        chatService.chatCompletionStream(chatCompletion, sseListener);
        
        return sseListener.getOutput().toString();
    }

    private ChatCompletion getChatCompletionBase() {
        WxModelConfigVo config = getConfig();
        return ChatCompletion.builder()
                .model(config.getModelType())
                .build();
    }

    private List<ChatMessage> buildChatMessageList(List<WxChatVo> chatVoList) {
        return chatVoList.stream().map(chatVo -> {
            if (chatVo.getRole().equals(Role.USER.getValue())) {
                return ChatMessage.withUser(chatVo.getContent());
            } else if (chatVo.getRole().equals(Role.ASSISTANT.getValue())) {
                return ChatMessage.withAssistant(chatVo.getContent());
            } else if (chatVo.getRole().equals(Role.SYSTEM.getValue())) {
                return ChatMessage.withSystem(chatVo.getContent());
            } else {
               throw new ServiceException("不支持的role类型");
            }
        }).collect(Collectors.toList());
    }
}
